from myImage.CImage import CImage
from myImage.CPatch import CPatch
from myImage.Helper import calcArea2

import math
import os
import numpy as np
import cv2

class CMerge:
    # 定义类的构造方法
    def __init__(self, columns, nums):
        self.columns = columns
        self.nums = nums

    # 获取检测结果文件，分子图片
    def getFiles(self, imgName, dir_dectect_res, iou, ratio, isLoadBoxFiles=True, rawDetects=[]):
        cols = self.columns
        nums = self.nums
        rows = nums / cols

        patchs = []  # ;(rows, cols) = CPatch();

        for j in range(self.nums):
            row = math.floor(j / cols)
            col = j % cols

            if len(patchs) < row + 1:
                patchs.append([])

            strSub = CImage.getString(j, cols)
            fileDetectRes = os.path.join(dir_dectect_res, '%s_%s.txt' % (imgName, strSub))
            if isLoadBoxFiles == True:
                objP = CPatch(fileDetectRes, iou, True, [])##读取子图片的检测结果bounding boxes
            else:
                objP = CPatch(fileDetectRes, iou, False, rawDetects[j])##根据

            objP.findEdgeBox()  ##; % % 查找边缘元素
            objP.supressEdges(ratio)  ##; % % 查找边缘元素

            patchs[row].append(objP)

        self.patchs = patchs
        return patchs

    # detects (class, prob, left, top, right, bottom, ..., ...)
    def getMatchPatches(self, boxes):
        patchs = self.patchs

        detects = np.empty((0, 8))

        cols = self.columns
        nums = self.nums
        rows = nums / cols

        for j in range(nums):
            row = math.floor(j / cols)
            col = j % cols

            box = boxes[j]
            left = box[0]
            top = box[1]

            patch = patchs[row][col]

            if len(patch.detects) == 0:
                continue

            offDet = np.array(patch.detects)

            offDet[:, 2] = offDet[:, 2] + left
            offDet[:, 3] = offDet[:, 3] + top
            offDet[:, 4] = offDet[:, 4] + left
            offDet[:, 5] = offDet[:, 5] + top

            detects = np.vstack((detects, offDet))

        return detects

    @staticmethod
    def NMSboxes(detects, threshold):
        if len(detects) <= 1:##只有一个box就不用抑制了
            return detects

        for label in range(3):  ## 分三个类进行NMS
            ind = np.where(detects[:, 0] == label)[0]
            if len(ind) <= 1:##少于一个没必要做
                continue

            bbox_array = detects[ind, 1:6]

            scores = bbox_array[:, 0]
            x1 = bbox_array[:, 1]
            y1 = bbox_array[:, 2]
            x2 = bbox_array[:, 3]
            y2 = bbox_array[:, 4]

            order = scores.argsort()[::-1]  # argsort函数返回的是数组值从小到大的索引值,[::-1]表示取反。即这里返回的是数组值从大到小的索引值
            areas = (x2 - x1 + 1) * (y2 - y1 + 1)  # 当前类所有框的面积(python会自动使用广播机制，相当于MATLAB中的.*即两矩阵对应元素相乘)；x1=3,x2=5,习惯上计算x方向长度就是x=3、4、5这三个像素，即5-3+1=3，而不是5-3=2，所以需要加1
            keep = []

            # 参考了https://zhuanlan.zhihu.com/p/40976906
            # 按confidence从高到低遍历bbx，移除所有与该矩形框的IoU值大于threshold的矩形框
            while order.size > 0:
                i = order[0]
                keep.append(i)  # 保留当前最大confidence对应的bbx索引
                # 获取所有与当前bbx的交集对应的左上角和右下角坐标，并计算IoU（注意这里是同时计算一个bbx与其他所有bbx的IoU）
                xx1 = np.maximum(x1[i], x1[order[1:]])  # 最大置信度的左上角坐标分别与剩余所有的框的左上角坐标进行比较，分别保存较大值；因此这里的xx1的维数应该是当前类的框的个数减1
                yy1 = np.maximum(y1[i], y1[order[1:]])
                xx2 = np.minimum(x2[i], x2[order[1:]])
                yy2 = np.minimum(y2[i], y2[order[1:]])
                inter = np.maximum(0.0, xx2 - xx1 + 1) * np.maximum(0.0, yy2 - yy1 + 1)
                iou = inter / (areas[i] + areas[order[1:]] - inter)  # 注意这里都是采用广播机制，同时计算了置信度最高的框与其余框的IoU
                inds = np.where(iou <= threshold)[0]  # 保留iou小于等于阙值的框的索引值
                order = order[inds + 1]  # 将order中的第inds+1处的值重新赋值给order；即更新保留下来的索引，加1是因为因为没有计算与自身的IOU，所以索引相差１，需要加上
                # bbox = bbox_array[keep]

            indind = ind[list(set(range(len(ind))).difference(set(keep)))]
            detects[indind, 0] = detects[indind, 0] + 128  ##表示抑制的box
        return detects

    @staticmethod
    def RefineBoxes(detects):##detects[,:1]已降序排列
        ins2 = np.where((detects[:, 0] == 2) | (detects[:, 6] == 1))[0]##6=1边缘元素
        ins1 = np.where((detects[:, 0] < 2) & (detects[:, 6] == 0))[0]

        # Step 1:把重叠的partial/边缘元素去除
        lend = len(ins2) ## 2023-02-23 修改，已排序
        for i in range(lend - 1):
            id = ins2[i]

            if detects[id, 0] >= 256:##增加
                continue
            seed = detects[id]

            for j in range(i + 1, lend):
                ii = ins2[j]

                if detects[ii, 0] >= 256:
                    continue

                part = detects[ii]

                (isBelong, _) = CMerge.isMostBelongto(part, seed, 0.5)
                if isBelong == True:
                    detects[ii, 0] = detects[ii, 0] + 256

        ## Step 2: 把重叠的seed去除
        lend = len(ins1)# 2023-02-24 新修改，修改seed_92右上角的双叶被(1,4)与(2,4)分别识别为seed与seed_1,导致互相+256
        for i in range(lend - 1):
            id = ins1[i]

            # if id == 107:
            #     kk = 0
            #     pass

            if detects[id, 0] >= 256:
                continue

            seedL = detects[id]  ##增加

            for j in range(i+1, lend):
                ii = ins1[j]

                if detects[ii, 0] >= 256:
                    continue

                part = detects[ii]

                # [isBelong, merge] = CMerge.isMostBelongto2(seedL, seed, 0.75, 0.15)
                # [isBelong, _] = CMerge.isMostBelongto2(part, seedL, 0.75, 0.15)
                [isBelong, _] = CMerge.isMostBelongto2(part, seedL, 0.65, 0.19)#bug: 2023-04-28 调试用

                if isBelong == True:
                    detects[ii, 0] = detects[ii, 0] + 256

        ## Step 3: 把 重叠在seed上的 边界元素与partial去掉 2023-02-24
        for id in ins1:
            if detects[id, 0] >= 256:
                continue

            seed = detects[id]

            for ii in ins2:
                if detects[ii, 0] >= 256:
                    continue

                part = detects[ii]

                [isBelong, _] = CMerge.isMostBelongto2(part, seed, 0.3, 0.8)###有问题
                if isBelong == True:
                    detects[ii, 0] = detects[ii, 0] + 256
        return detects

    @staticmethod
    def isMostBelongto(setPart, setBig, thresh):
        isBelong = False
        merge = []

        [isInside, areaLap] = CMerge.isInsideBig(setPart, setBig)
        if isInside == False:
            return isBelong, merge

        areaPart = (setPart[5] - setPart[3] + 1) * (setPart[4] - setPart[2] + 1)
        areaBig  = (setBig[5] - setBig[3] + 1) * (setBig[4] - setBig[2] + 1)

        thresh1 = thresh * 0.9 ##% % 针对partial
        thresh2 = thresh ##% % 针对边界元素

        label = int(setPart[0])
        if (areaLap / areaPart < thresh1) and (label & 2 == True):
            return isBelong, merge
        elif (areaLap / areaPart < thresh2) and (label & 2 == False):
            return isBelong, merge

        isBelong = True
        merge = CMerge.CalcMerge(setPart, setBig)
        return isBelong, merge

    @staticmethod
    def isMostBelongto2(setPart, setBig, thresh1, thresh2):
        isBelong = False
        merge = []

        [isInside, areaLap] = CMerge.isInsideBig(setPart, setBig)
        if isInside == False:
            return isBelong, merge

        areaPart = (setPart[5] - setPart[3] + 1) * (setPart[4] - setPart[2] + 1)##预先算好面积，下次改！！！
        areaBig = (setBig[5] - setBig[3] + 1) * (setBig[4] - setBig[2] + 1)

        if areaLap / areaPart < thresh1:
            return isBelong, merge

        if (areaPart - areaLap) / areaBig > thresh2:
            print('CMerge(isMostBelongto2):----有一个异常！！')
            return isBelong, merge

        isBelong = True
        merge = CMerge.CalcMerge(setPart, setBig)
        return isBelong, merge

    @staticmethod
    def isInsideBig(setPart, setBig):##表示重叠区域 是否在Big种子内部
        leftMax = max(setPart[2], setBig[2])
        topMax = max(setPart[3], setBig[3])
        rightMin = min(setPart[4], setBig[4])
        bottomMin = min(setPart[5], setBig[5])

        areaLap = (bottomMin - topMax + 1) * (rightMin - leftMax + 1)
        if bottomMin < topMax or rightMin < leftMax:
            isBelongto = False
        else:
            isBelongto = True
        return isBelongto, areaLap

    @staticmethod
    def CalcMerge(setPart, setBig):
        leftMin = min(setPart[2], setBig[2])
        topMin = min(setPart[3], setBig[3])
        rightMax = max(setPart[4], setBig[4])
        bottomMax = max(setPart[5], setBig[5])

        merge = [leftMin, topMin, rightMax, bottomMax]
        return merge

    @staticmethod
    def SaveImageAndBoxes(img, fileSaveDrawBox, detects, isDraw128, ratio=50): #ratio:00/100
        for array in detects:
            label = array[0]

            box = np.array(array[2:6], dtype=int)

            if label == 0:
                cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 2)
            elif label == 1:
                cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (255, 0, 0), 2)
            elif label == 2:
                cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 0, 0), 2)
            elif label >= 128:
                if isDraw128 == False:
                    continue

                label2 = label - 128
                if label2 == 0:
                    cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (128, 128, 0), 2)
                elif label2 == 1:
                    cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (128, 0, 128), 2)
                elif label2 == 2:
                    cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 128, 128), 2)
                else:
                    cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (128, 128, 128), 2)
            else:
                print('error....')

        print('CMerge(SaveImageAndBoxes):', fileSaveDrawBox)
        params = [cv2.IMWRITE_JPEG_QUALITY, ratio]
        cv2.imencode('.jpg', img, params)[1].tofile(fileSaveDrawBox)##不能用imwrite
        # cv2.imencode('.jpg', img)[1].tofile(fileSaveDrawBox)##不能用imwrite
        return img

    @staticmethod
    def DrawLines(img):##画网格线
        height, width, _ = img.shape

        ww = 650
        rows = list(range(0, height - ww, ww))
        cols = list(range(0, width  - ww, ww))

        i = 1
        for row in rows:
            cv2.rectangle(img, (0, row), (cols[-1] + ww, row + ww), (100, 100, 128), 2)
            cv2.putText(img, '%d' % i, (row + 250, row + 250),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        6, (255, 0, 0, 128), 2) #增加显示数字
            i = i + 1

        for col in cols:
            cv2.rectangle(img, (col, 0), (col + ww, rows[-1] + ww), (100, 100, 128), 2)
        return img


    @staticmethod
    def StatSeeds(detects):
        seeds = np.sum(detects[:, 0] == 0)
        seed_1 = np.sum(detects[:, 0] == 1)
        partial = np.sum(detects[:, 0] == 2)
        data = [seeds, seed_1, partial, seeds + seed_1]
        return data

    @staticmethod
    def SaveDetectsToFile(fileName, detects):
        with open(fileName, 'w') as f:
            for detect in detects:
                str =  [f'{i:.3f}\t'  for i in detect]
                str2 = ''.join(str) + '\n'
                f.write((str2))